{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of W10_Tutorial2_Instructor.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nn57geGyp2hO"
      },
      "source": [
        "# CIS-522 Week 10 Part 2\n",
        "# Introduction to Transformers, BERT, and Language Models\n",
        "\n",
        "**Instructor:** Lyle Ungar\n",
        "\n",
        "**Content Creators:** Sanjeevini  Ganni, Ameet Rahane, Byron Galbraith\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ymkujLBlqM6F"
      },
      "source": [
        "##Tutorial Objectives\n",
        "\n",
        "\n",
        "(1) Recognize NLP tasks: IR/search,  Question Answering/text completion, MT \\\\\n",
        "(2) Understand distributional similarity on words and context including Context-oblivious embeddings (word2vec, glove, fastText) and multilingual embeddings  \\\\\n",
        "(3) Attention  \\\\\n",
        "(4) Context-sensitive embeddings: BERT and transformers: masking and self-attention \\\\\n",
        "(5) The many flavors of BERT:  RoBERTa and DistilBERT \\\\\n",
        "(6) Fine-tuning language embeddings \\\\\n",
        "(7) Explaining NLP models \\\\\n",
        "(8) Big language models: GPT-3 and friends  \\\\\n",
        "(9) Societal: Bias in language embeddings \\\\\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VC2Bqp7M12U1"
      },
      "source": [
        "---\n",
        "## Preface\n",
        "We recommend saving this notebook in your Google Drive (`File -> Save a copy in Drive`)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LsvDRvFG15Rw",
        "cellView": "form"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "my_pennkey = '' #@param {type:\"string\"}\n",
        "my_pod = 'Select' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion','astute-jellyfish', 'quantum-herring']\n",
        "\n",
        "# start timing\n",
        "import time\n",
        "try:t0;\n",
        "except NameError: t0 = time.time()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2_oC9dNZqQu-"
      },
      "source": [
        "##Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0RcjygJ2CRSG",
        "cellView": "form"
      },
      "source": [
        "#@title Install\n",
        "!pip install torchtext==0.4.0\n",
        "!pip install transformers\n",
        "# !git clone https://github.com/facebookresearch/fastText.git\n",
        "# %cd fastText\n",
        "# !pip install . \n",
        "# %cd .."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "52U7jyR4pngL",
        "cellView": "form"
      },
      "source": [
        "#@title Imports and Seed\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.cm as cm\n",
        "% matplotlib inline\n",
        "import re\n",
        "from IPython.display import Image\n",
        "import os\n",
        "from tqdm import tqdm_notebook as tqdm\n",
        "import sys\n",
        "import random\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "import torch.optim as optim\n",
        "from torch.autograd import Variable\n",
        "from torchtext import data, datasets\n",
        "from torchtext.vocab import Vectors, FastText\n",
        "# import fasttext\n",
        "import requests\n",
        "import zipfile\n",
        "# import nltk\n",
        "# nltk.download('punkt')\n",
        "# from nltk.tokenize import word_tokenize\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
        "import urllib\n",
        "import csv\n",
        "from scipy.special import softmax\n",
        "\n",
        "seed = 522\n",
        "random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed_all(seed)\n",
        "torch.cuda.manual_seed(seed)\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l82sJ8BNmty-",
        "cellView": "form"
      },
      "source": [
        "# @title Figure Settings\n",
        "import ipywidgets as widgets\n",
        "%matplotlib inline \n",
        "fig_w, fig_h = (8, 6)\n",
        "plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "SMALL_SIZE = 12\n",
        "\n",
        "\n",
        "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/\"\n",
        "              \"course-content/master/nma.mplstyle\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ik8be94eRLPt",
        "cellView": "form"
      },
      "source": [
        "#@title Helper functions\n",
        "def cosine_similarity(vec_a, vec_b):\n",
        "        \"\"\"Compute cosine similarity between vec_a and vec_b\"\"\"\n",
        "        return np.dot(vec_a, vec_b) / \\\n",
        "            (np.linalg.norm(vec_a) * np.linalg.norm(vec_b))\n",
        "def tokenize(sentences):\n",
        "  #Tokenize the sentence\n",
        "  #from nltk.tokenize library use word_tokenize\n",
        "  token = word_tokenize(sentences)\n",
        "  \n",
        "  return token\n",
        "\n",
        "def plot_train_val(x, train, val, train_label, val_label, title):\n",
        "  plt.plot(x, train, label=train_label)\n",
        "  plt.plot(x, val, label=val_label)\n",
        "  plt.legend(loc='lower right')\n",
        "  plt.xlabel('epoch')\n",
        "  plt.title(title)\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bi_7aepGAZkf"
      },
      "source": [
        "###Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xCZ56QJ96aMW",
        "cellView": "form"
      },
      "source": [
        "#@title Load Data\n",
        "\n",
        "def load_dataset(emb_vectors, sentence_length = 50):\n",
        "    TEXT = data.Field(sequential=True, tokenize=tokenize, lower=True, include_lengths=True, batch_first=True, fix_length=sentence_length)\n",
        "    LABEL = data.LabelField(dtype=torch.float)\n",
        "\n",
        "    train_data, test_data = datasets.IMDB.splits(TEXT, LABEL)\n",
        "\n",
        "    TEXT.build_vocab(train_data, vectors=emb_vectors)\n",
        "    LABEL.build_vocab(train_data)\n",
        "\n",
        "    train_data, valid_data = train_data.split(split_ratio=0.7, random_state = random.seed(seed))\n",
        "    train_iter, valid_iter, test_iter = data.BucketIterator.splits((train_data, valid_data, test_data), batch_size=32, sort_key=lambda x: len(x.text), repeat=False, shuffle=True)\n",
        "    vocab_size = len(TEXT.vocab)\n",
        "\n",
        "    return TEXT, vocab_size, train_iter, valid_iter, test_iter\n",
        "\n",
        "\n",
        "def download_file_from_google_drive(id, destination):\n",
        "    URL = \"https://docs.google.com/uc?export=download\"\n",
        "\n",
        "    session = requests.Session()\n",
        "\n",
        "    response = session.get(URL, params = { 'id' : id }, stream = True)\n",
        "    token = get_confirm_token(response)\n",
        "\n",
        "    if token:\n",
        "        params = { 'id' : id, 'confirm' : token }\n",
        "        response = session.get(URL, params = params, stream = True)\n",
        "\n",
        "    save_response_content(response, destination)    \n",
        "\n",
        "def get_confirm_token(response):\n",
        "    for key, value in response.cookies.items():\n",
        "        if key.startswith('download_warning'):\n",
        "            return value\n",
        "\n",
        "    return None\n",
        "\n",
        "def save_response_content(response, destination):\n",
        "    CHUNK_SIZE = 32768\n",
        "\n",
        "    with open(destination, \"wb\") as f:\n",
        "        for chunk in response.iter_content(CHUNK_SIZE):\n",
        "            if chunk: # filter out keep-alive new chunks\n",
        "                f.write(chunk)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IWxI-yHszqqs"
      },
      "source": [
        "---\n",
        "##Section 1: Transformers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8J7i4TZQAIix",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Self-attention\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"g860drKesIw\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HJhcL3PivRyM"
      },
      "source": [
        "Transformers! Like CNNs and LSTMs, this base model architecture has been the foundation of many very successful models such as BERT and friends. It has seen tremendous use in the NLP space, and there have been effforts to extend it to the audio and visual domains as well. The original paper, [Attention Is All You Need (Vaswani et al. 2017)](https://arxiv.org/abs/1706.03762), is very readable and you are encouraged to take a look.\n",
        "\n",
        "The Transformer model is fundamentally an encoder-decoder model that operates on sequences of tokens. Both the encoder and decoder components are composed of stacks of submodules that use **only** attention mechanisms and linear weights to learn (there are no CNNs or RNNs). The architecture schematic looks like the following:\n",
        "\n",
        "![transformer architecture](https://d2l.ai/_images/transformer.svg)\n",
        "\n",
        "In the rest of this section we will be going over the various building blocks that go into Transformers. The goal here is not to train anything (that is left for a homework assignment). Rather, the emphasis is on understanding what all the pieces do and how they fit together.\n",
        "\n",
        "*Note:* Many of the images in this section are taken from Dive Into Deep Learning's chapter on [Attention Mechanisms](https://d2l.ai/chapter_attention-mechanisms/index.html). You are encouraged to check that out for additional details and implementations.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YjdBlVNuT96x"
      },
      "source": [
        "### Self-Attention\n",
        "\n",
        "Transformers make use of something called self-attention as a critical component to the entire operation. What does that mean in the context of attention mechanisms? If you recall, attention mechanisms in machine learning have three components:\n",
        "\n",
        " - the values V (the things you perceive i.e. model inputs)\n",
        " - the query Q (the thing you want to attend to)\n",
        " - the keys K (a mapping between queries and values)\n",
        " \n",
        "Generally the number and dimensionality of queries and values can all be different. In self-attention, the queries, keys, and values are all drawn from the same set of inputs. In other words, we don't need to specify anything about what and how queries and keys are formed, as they come straight from the data just like the values!\n",
        "\n",
        "Take a minute and look at this article from the last pod session. It has detailed graphical explanation on how to calculate attention scores. \n",
        "https://towardsdatascience.com/illustrated-self-attention-2d627e33b20a\n",
        "\n",
        "Ok, so we know that our queries, keys, and values come from our input sequence, but which attention mechanism should we use?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "djMSZnkXkeDU"
      },
      "source": [
        "### Masked Scaled Dot Product Attention"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hvk5zc2akcM0",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Masking\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"VtaGIp_9j1w\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TLCjYZvPR4-A"
      },
      "source": [
        "Masking is used to hide values from being attended to, including \\\n",
        "(a) words that are hidden for self-supervised learning, \\\n",
        "(b) padding tokens, and \\\n",
        "(c) in the seq2seq tasks in the original Transformers paper, everything that came after the next token in the output training (to enforce autoregressive behavior). \n",
        "\n",
        "BERT, on the other hand, masks out individual tokens within the input sequence and everything but those tokens on the output side. This provides a more bidirectional-style embedding learning strategy. \n",
        "\n",
        "This is accomplished by setting every element we want to mask to $-\\infty$ before applying softmax, which has the effect of giving that element a probability of 0.\n",
        "\n",
        "We've provided a masked softmax function below, which assumes a binary matrix of size (batch_size, n_tokens) where a value of 0 indicates that token should be masked.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_kodsdE5dedu"
      },
      "source": [
        "def masked_softmax(x, mask):\n",
        "  \"\"\" Applies softmax on a masked version of the input.\n",
        "  Args:\n",
        "    x (n_batch, n_tokens, t_tokens): - the scaled dot product of Q and K\n",
        "    mask (n_batch, n_tokens): - binary mask, all values = 0 will be set to -inf\n",
        "  Returns:\n",
        "    (n_batch, n_tokens, n_tokens): the result of applying softmax along the last\n",
        "                                   dimension of the masked input.\n",
        "  \"\"\"\n",
        "  return F.softmax(x.masked_fill_(mask.unsqueeze(1) == 0, float('-inf')), dim=-1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nKRJULI3CM_b"
      },
      "source": [
        "#### Exercise 1: Masked Scaled Dot Product Attention Module\n",
        "In this exercise you will implement the forward method of a PyTorch module for computing the masked scaled dot product attention function. This is represented by the following equation:\n",
        "\n",
        "$$\n",
        "\\alpha(Q, K, V, M) = \\mathrm{masked\\ softmax} \\left( \\frac{QK^T}{\\sqrt d_k}, M \\right) V\n",
        "$$\n",
        "\n",
        "where $Q$ is the query tensor, $K$ is the key tensor, $V$ is the value tensor, $M$ is the masking tensor, and $d_k$ is the dimensionality of our embeddings.\n",
        "\n",
        "PyTorch provides us with the very useful [`torch.bmm`](https://pytorch.org/docs/stable/generated/torch.bmm.html) function to compute matrix multiplication over batches while preserving the batch dimension. You will also want to make use of the [`torch.transpose`](https://pytorch.org/docs/stable/generated/torch.transpose.html#torch.transpose) function to get the transpose of the keys tensor while also preserving the batch dimension.\n",
        "\n",
        "To calculate the masked scaled dot product attention, we will first get the dot product from the query and key. We divide the dot product by the square root of the embedding length to scale it. Then we apply the masked softmax function to get the scores. By multiplying the scores and the values we get the masked scaled dot product attention.\n",
        "\n",
        "*NOTE:* Dropout is normally applied to the `scaled_dot_product` quantity before softmax is applied during training. However, in the interests of clarity, we are omitting it here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KguOtPV5ClSA"
      },
      "source": [
        "class ScaledDotProductAttention(nn.Module):\n",
        "  def __init__(self, embed_dim):    \n",
        "    super().__init__()\n",
        "    self.embed_dim = embed_dim\n",
        "\n",
        "  def forward(self, queries, keys, values, mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      queries (n_batch, n_tokens, embed_dim): queries (Q) tensor\n",
        "      keys (n_batch, n_tokens, embed_dim): keys (K) tensor\n",
        "      values (n_batch, n_tokens, embed_dim): values (V) tensor\n",
        "      mask (n_batch, n_tokens, embed_dim): binary mask (M) tensor\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): scaled dot product attention tensor\n",
        "    \"\"\"\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"ScaledDotProductAttention\")\n",
        "    ####################################################################\n",
        "    # Compute the dot product of the queries and keys, then divide by the square\n",
        "    # root of the embedding size\n",
        "    scaled_dot_product = ...\n",
        "    # Next perform the masked softmax function on this product and supplied mask\n",
        "    masked_softmax_scores = ...\n",
        "    # Finally compute and return the dot product of the masked output and the values\n",
        "    attention = ...\n",
        "    return attention\n",
        "\n",
        "# Uncomment below to test your module\n",
        "# torch.manual_seed(522)\n",
        "# batch_size, n_tokens, embed_dim = 1, 3, 4\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# attention = ScaledDotProductAttention(embed_dim)\n",
        "# mask = torch.ones((batch_size, n_tokens))\n",
        "# print(attention(tokens, tokens, tokens, mask))\n",
        "# mask[0, 2:] = 0\n",
        "# print(attention(tokens, tokens, tokens, mask))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ENqM-9EevHr-"
      },
      "source": [
        "# to_remove solution\n",
        "class ScaledDotProductAttention(nn.Module):\n",
        "  def __init__(self, embed_dim):\n",
        "    super().__init__()\n",
        "    self.embed_dim = embed_dim\n",
        "\n",
        "  def forward(self, queries, keys, values, mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      queries (n_batch, n_tokens, embed_dim): queries (Q) tensor\n",
        "      keys (n_batch, n_tokens, embed_dim): keys (K) tensor\n",
        "      values (n_batch, n_tokens, embed_dim): values (V) tensor\n",
        "      mask (n_batch, n_tokens): binary mask tensor\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): scaled dot product attention tensor\n",
        "    \"\"\"\n",
        "    scaled_dot_product = torch.bmm(queries, torch.transpose(keys, 1, 2)) / np.sqrt(self.embed_dim)\n",
        "    masked_softmax_scores = masked_softmax(scaled_dot_product, mask)\n",
        "    attention = torch.bmm(masked_softmax_scores, values)\n",
        "    return attention\n",
        "\n",
        "# torch.manual_seed(522)\n",
        "# batch_size, n_tokens, embed_dim = 1, 3, 4\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# attention = ScaledDotProductAttention(embed_dim)\n",
        "# mask = torch.ones((batch_size, n_tokens))\n",
        "# print(attention(tokens, tokens, tokens, mask))\n",
        "# mask[0, 2:] = 0\n",
        "# print(attention(tokens, tokens, tokens, mask))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCxrHyzvhHCv"
      },
      "source": [
        "If done correctly, you should see something like the following (numbers may vary but shape should be the same):\n",
        "```\n",
        "tensor([[[-0.2151, -0.0403,  0.9237, -1.6000],\n",
        "         [ 0.4216,  1.3972,  1.3613, -0.0161],\n",
        "         [ 0.0862,  0.7809,  1.1435, -0.8234]]])\n",
        "tensor([[[-0.2688, -0.3399,  0.8707, -1.7680],\n",
        "         [ 0.5448,  1.5134,  1.4313,  0.2591],\n",
        "         [ 0.1679,  0.6550,  1.1717, -0.6798]]])\n",
        "```\n",
        "[*Click for solution*](https://github.com/CIS-522/course-content/blob/main/tutorials/W10_NLP/solutions/W10_Tutorial2_Solution_Ex01.py)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2fVCkImZNxq"
      },
      "source": [
        "Self-attention is great, but it has two shortcomings:\n",
        "\n",
        "1. It doesn't let us specify or control what gets attended to and thus will converge on only one strategy due to averaging. (Like short-term attention or long-term attention).\n",
        "\n",
        "2. There is no implicit ordering or notion of relative position of the input tokens to each other unlike in RNNs and ConvNets.\n",
        "\n",
        "We know things about natural language, such as that word order matters and there are many different grammatical and syntactic features that imbue useful meaning. How do we overcome this?\n",
        "\n",
        "First, let's address the attention strategy problem. One answer to only having a single attention strategy is to have many!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Onssr49pLFiK"
      },
      "source": [
        "### Multi-Head Attention\n",
        "\n",
        "In practice, given the same set of queries, keys, and values we may want our model to combine knowledge from different behaviors of the same attention mechanism, such as capturing dependencies of various ranges (e.g., shorter-range vs. longer-range) within a sequence. Thus, it may be beneficial to allow our attention mechanism to jointly use different representation subspaces of queries, keys, and values.\n",
        "\n",
        "Multi-head attention mechanism is employed by Transformers to concurrently learn multiple different attention strategies or \"heads.\" This is accomplished by passing each of the queries, keys, and values through single, fully-connected linear layers. Attention training is then conducted on all splits, which then get joined together at the end and passed through a linear layer to achieve the final output.\n",
        "\n",
        "![multi-head attention diagram](https://d2l.ai/_images/multi-head-attention.svg)\n",
        "\n",
        "Now, to avoid poor scaling performance with each attention head we add, we can take advantage of the fact that we only need to compute matrix multiplications. By effectively making the dimensionality of the newly made query, key, and value heads equal to the original embedding dimension cleanly divided by the number of heads, we can keep the heads strided in one tensor and thus compute the attention scores of all heads in a single call.\n",
        "\n",
        "The methods to shuffle the data around for the input values (queries, keys, values) and then to unshuffle it for the output are provided below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hWdhsAo_3qa6"
      },
      "source": [
        "def mha_transform_input(x, n_heads, head_dim):\n",
        "  \"\"\" Restructure the input tensors to compute the heads in parallel\n",
        "  Requires that head_dim = embed_dim / n_heads\n",
        "  Args:\n",
        "    x (n_batch, n_tokens, embed_dim): input tensor, one of queries, keys, or values\n",
        "    n_heads (int): the number of attention heads\n",
        "    head_dim (int): the dimensionality of each head\n",
        "  Returns:\n",
        "    (n_batch*n_heads, n_tokens, head_dim): 3D Tensor containing all the input heads\n",
        "  \"\"\"\n",
        "  n_batch, n_tokens, _ = x.shape    \n",
        "  x = x.reshape((n_batch, n_tokens, n_heads, head_dim))\n",
        "  x = x.permute(0, 2, 1, 3)\n",
        "  return x.reshape((n_batch * n_heads, n_tokens, head_dim))\n",
        "  \n",
        "def mha_transform_output(x, n_heads, head_dim):\n",
        "  \"\"\" Restructures the output back to the original format \n",
        "  Args:\n",
        "    x (n_bacth*n_heads, n_tokens, head_dim): multi-head representation tensor\n",
        "    n_heads (int): the number of attention heads\n",
        "    head_dim (int): the dimensionality of each head\n",
        "  Returns:\n",
        "    (n_batch, n_tokens, embed_dim): 3D Tensor containing all the input heads\n",
        "  \"\"\"\n",
        "  n_concat, n_tokens, _ = x.shape\n",
        "  n_batch = n_concat // n_heads\n",
        "  x = x.reshape((n_batch, n_heads, n_tokens, head_dim))\n",
        "  x = x.permute(0, 2, 1, 3)\n",
        "  return x.reshape((n_batch, n_tokens, n_heads * head_dim))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-_rWeVEtKgdk"
      },
      "source": [
        "#### Exercise 2: Multi-Head Attention Module\n",
        "In this exercise you will implement the the forward method of a PyTorch module for handling the multi-head attention mechanism. Each of the Q, K, and V inputs need to be run through their corresponding linear layers and then transformed using `mha_transform_input`. You then pass these to our scaled dot product attention module, transform that output back using `mha_transform_output`, and then run that though the corresponding output linear layer.\n",
        "\n",
        "*NOTE:* In the original Transformers paper, the linear layers were just weight matrices with no bias term which is reproduced here by using `Linear` layers and setting bias to False."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_Ohoi8GWGN8l"
      },
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "  def __init__(self, n_heads, embed_dim):\n",
        "    super().__init__()\n",
        "    self.n_heads = n_heads\n",
        "    self.head_dim = embed_dim // n_heads\n",
        "    \n",
        "    self.attention = ScaledDotProductAttention(embed_dim)\n",
        "    self.query_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "    self.key_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "    self.value_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "    self.out_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "  \n",
        "  def forward(self, queries, keys, values, mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      queries (n_batch, n_tokens, embed_dim): queries (Q) tensor\n",
        "      keys (n_batch, n_tokens, embed_dim): keys (K) tensor\n",
        "      values (n_batch, n_tokens, embed_dim): values (V) tensor\n",
        "      mask (n_batch, n_tokens): binary mask tensor\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): multi-head attention tensor\n",
        "    \"\"\"\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"MultiHeadAttention\")\n",
        "    ####################################################################\n",
        "    # Pass the queries through their linear layer, then apply the mha_transform_input function\n",
        "    q_heads = ...\n",
        "    # Pass the keys through their linear layer, then apply the mha_transform_input function\n",
        "    k_heads = ...\n",
        "    # Pass the values through their linear layer, then apply the mha_transform_input function\n",
        "    v_heads = ...\n",
        "\n",
        "    # Compute the scaled dot product attention on the transformed q, k, and v\n",
        "    # attention heads with the provided mask\n",
        "    attention_heads = ...\n",
        "    # Apply the mha_transform_output function to the attention heads, then pass\n",
        "    # this through the output linear layer\n",
        "    attention = ...\n",
        "    return attention\n",
        "\n",
        "# Uncomment below to test your module\n",
        "# torch.manual_seed(522)\n",
        "# n_heads, batch_size, n_tokens, embed_dim = 2, 1, 3, 4\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# mask = torch.ones((batch_size, n_tokens))\n",
        "# attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "# attention(tokens, tokens, tokens, mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q2uPKOySvRO1"
      },
      "source": [
        "# to_remove solution\n",
        "class MultiHeadAttention(nn.Module):\n",
        "  def __init__(self, n_heads, embed_dim):\n",
        "    super().__init__()\n",
        "    self.n_heads = n_heads\n",
        "    self.head_dim = embed_dim // n_heads\n",
        "    \n",
        "    self.attention = ScaledDotProductAttention(embed_dim)\n",
        "    self.query_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "    self.key_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "    self.value_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "    self.out_fc = nn.Linear(embed_dim, embed_dim, bias=False)\n",
        "  \n",
        "  def forward(self, queries, keys, values, mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      queries (n_batch, n_tokens, embed_dim): queries (Q) tensor\n",
        "      keys (n_batch, n_tokens, embed_dim): keys (K) tensor\n",
        "      values (n_batch, n_tokens, embed_dim): values (V) tensor\n",
        "      mask (n_batch, n_tokens): binary mask tensor\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): multi-head attention tensor\n",
        "    \"\"\"\n",
        "    q_heads = mha_transform_input(self.query_fc(queries), self.n_heads, self.head_dim)\n",
        "    k_heads = mha_transform_input(self.key_fc(keys), self.n_heads, self.head_dim)\n",
        "    v_heads = mha_transform_input(self.value_fc(values), self.n_heads, self.head_dim)\n",
        "\n",
        "    attention_heads = self.attention(q_heads, k_heads, v_heads, mask)\n",
        "    attention = self.out_fc(mha_transform_output(attention_heads, self.n_heads, self.head_dim))\n",
        "    return attention\n",
        " \n",
        "# torch.manual_seed(522)\n",
        "# n_heads, batch_size, n_tokens, embed_dim = 2, 1, 3, 4\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# mask = torch.ones((batch_size, n_tokens))\n",
        "# attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "# attention(tokens, tokens, tokens, mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gshY1bdKx-IM"
      },
      "source": [
        "If done correctly, you should see something like the following (numbers may vary but shape should be the same):\n",
        "```\n",
        "tensor([[[ 0.1191,  0.3588,  0.2972, -0.2594],\n",
        "         [ 0.1204,  0.3102,  0.2904, -0.2539],\n",
        "         [ 0.1216,  0.3362,  0.2921, -0.2603]]], grad_fn=<UnsafeViewBackward>)\n",
        "```\n",
        "[*Click for solution*](https://github.com/CIS-522/course-content/blob/main/tutorials/W10_NLP/solutions/W10_Tutorial2_Solution_Ex02.py)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JcVRwWnKf3ZN"
      },
      "source": [
        "So we have a solution for enabling multiple different attention strategies to be applied. But how about knowing where each token is in the sequence? Unlike RNNs that recurrently process tokens of a sequence one by one, self-attention ditches sequential operations in favor of parallel computation. Well, what if we just explicitly added some piece of information to the input representation that encodes each token's position?\n",
        "\n",
        "Positional encodings can be either learned or fixed. While you can likely imagine ways to do this, Transformers use fixed positional encoding based on sine and cosine functions:\n",
        "\n",
        "$$\n",
        "p_{i,2j} = \\sin \\left( \\frac{i}{10000^{2j/d}} \\right) \\\\\n",
        "p_{i,2j+1} = \\cos \\left( \\frac{i}{10000^{2j/d}} \\right)\n",
        "$$\n",
        "\n",
        "where $i$ and $j$ are iterated over the rows (tokens) and columns (embedding dimensions), respectively. This likely seems strange at first, but it has the neat effect of\n",
        "1. providing unique values across the matrix elements\n",
        "2. uses float values which easily add to the input embedded tokens\n",
        "\n",
        "We can see an example of what this looks like when plotted for a few columns below:\n",
        "\n",
        "![positional embedding](https://d2l.ai/_images/output_self-attention-and-positional-encoding_d76d5a_31_0.svg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BpM1VQmqLXQT"
      },
      "source": [
        "### Exercise 3: Positional Encoding Module\n",
        "In this exercise you will create the forward method for a PyTorch module that will add positional embeddings to an input batch of tokens. The position embedding values are already computed and cached for you.\n",
        "\n",
        "*NOTE:* Dropout is normally applied to the output of this module during training, but we have omitted it for clarity."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i8cAATbDAiA4",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Positional Encoding\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"FoRWkEAJDtg\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "elotwfaIIM1o"
      },
      "source": [
        "class PositionalEncoder(nn.Module):\n",
        "  def __init__(self, embed_dim, max_len=1000):\n",
        "    super().__init__()\n",
        "    self.position_embedding = torch.zeros((1, max_len, embed_dim))\n",
        "    i = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1)\n",
        "    j2 = torch.arange(0, embed_dim, step=2, dtype=torch.float32)\n",
        "    x = i / torch.pow(10000, j2 / embed_dim)\n",
        "    self.position_embedding[..., 0::2] = torch.sin(x)\n",
        "    self.position_embedding[..., 1::2] = torch.cos(x)        \n",
        "\n",
        "  def forward(self, x):\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"PositionalEncoder\")\n",
        "    ####################################################################\n",
        "    # Add the cached positional encodings to the input\n",
        "    x_plus_p = ...\n",
        "    return x_plus_p\n",
        "\n",
        "# Uncomment below to test your module\n",
        "# n_tokens, embed_dim = 10, 4\n",
        "# pos_enc = PositionalEncoder(embed_dim) \n",
        "# p = pos_enc(torch.zeros((1, n_tokens, embed_dim)))\n",
        "# plt.imshow(p.squeeze())\n",
        "# p"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IoxUJUk8vXf2"
      },
      "source": [
        "# to_remove solution\n",
        "class PositionalEncoder(nn.Module):\n",
        "  def __init__(self, embed_dim, max_len=1000):\n",
        "    super().__init__()\n",
        "    self.position_embedding = torch.zeros((1, max_len, embed_dim))\n",
        "    i = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1)\n",
        "    j2 = torch.arange(0, embed_dim, step=2, dtype=torch.float32)\n",
        "    x = i / torch.pow(10000, j2 / embed_dim)\n",
        "    self.position_embedding[..., 0::2] = torch.sin(x)\n",
        "    self.position_embedding[..., 1::2] = torch.cos(x)        \n",
        "\n",
        "  def forward(self, x):\n",
        "    x_plus_p = x + self.position_embedding[:, :x.shape[1]]\n",
        "    return x_plus_p\n",
        "  \n",
        "# with plt.xkcd():\n",
        "#   n_tokens, embed_dim = 40, 40\n",
        "#   pos_enc = PositionalEncoder(embed_dim) \n",
        "#   p = pos_enc(torch.zeros((1, n_tokens, embed_dim)))\n",
        "#   plt.imshow(p.squeeze())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T34SUKdayCvt"
      },
      "source": [
        "You should see a plot visualizing the different values, as well as the actual positional output values:\n",
        "```\n",
        "tensor([[[ 0.0000,  1.0000,  0.0000,  1.0000],\n",
        "         [ 0.8415,  0.5403,  0.0100,  0.9999],\n",
        "         [ 0.9093, -0.4161,  0.0200,  0.9998],\n",
        "         [ 0.1411, -0.9900,  0.0300,  0.9996],\n",
        "         [-0.7568, -0.6536,  0.0400,  0.9992],\n",
        "         [-0.9589,  0.2837,  0.0500,  0.9988],\n",
        "         [-0.2794,  0.9602,  0.0600,  0.9982],\n",
        "         [ 0.6570,  0.7539,  0.0699,  0.9976],\n",
        "         [ 0.9894, -0.1455,  0.0799,  0.9968],\n",
        "         [ 0.4121, -0.9111,  0.0899,  0.9960]]])\n",
        "```\n",
        "[*Click for solution*](https://github.com/CIS-522/course-content/blob/main/tutorials/W10_NLP/solutions/W10_Tutorial2_Solution_Ex03.py)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4N7_j4uihthC"
      },
      "source": [
        "### The Encoder\n",
        "\n",
        "We now have almost everything we need to assemble the full Transformer network. There are just two more modules we need to quickly discuss, and then we will get to putting them all together.\n",
        "\n",
        "Transformer architecture for reference:\n",
        "<div>\n",
        "<img src=\"https://d2l.ai/_images/transformer.svg\" width=\"275\"/>\n",
        "</div>\n",
        "\n",
        "First, there is the residual layer norm that appears after every other component. In all cases, this takes the output of the previous component, sums it with the input to that component (the residual connection), and then normalizes the result across the layer.\n",
        "\n",
        "Second is the positionwise feed forward network that appears after the attention components. It is a two layer fully connected module with a ReLU activation in between. \n",
        "\n",
        "These are provided below. Note that dropout would normally be applied in various places in these modules during training, but we have omitted it for clarity."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YAF3g4h51_dA"
      },
      "source": [
        "class ResidualNorm(nn.Module):\n",
        "  def __init__(self, embed_dim):\n",
        "    super().__init__()\n",
        "    self.norm = nn.LayerNorm(embed_dim)\n",
        "  \n",
        "  def forward(self, x, residual):\n",
        "    return self.norm(x + residual)\n",
        "\n",
        "class Feedforward(nn.Module):\n",
        "  def __init__(self, embed_dim, hidden_dim):\n",
        "    super().__init__()\n",
        "    self.fc1 = nn.Linear(embed_dim, hidden_dim)    \n",
        "    self.fc2 = nn.Linear(hidden_dim, embed_dim)\n",
        "  \n",
        "  def forward(self, x):\n",
        "    return self.fc2(F.relu(self.fc1(x)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0DSh0XVD1_Nl"
      },
      "source": [
        "Now that we have all the modules we need, we can begin assembling the bigger modules. First we will look at the Encoder Block. The actual encoder will be made up of some number of these stacked together."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkOaVXrnU-43"
      },
      "source": [
        "#### Exercise 4: Encoder Block Module\n",
        "In this exercise you will create the forward method of the PyTorch module representing the Encoder Block of the Transformer. The Encoder Block has the following architecture:\n",
        "1. a multi-head attention module using self-attention\n",
        "2. 1st residual layer norm\n",
        "3. feed-forward model\n",
        "4. 2nd residual layer norm\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NwIxWxc3S7CT"
      },
      "source": [
        "class EncoderBlock(nn.Module):\n",
        "  def __init__(self, n_heads, embed_dim, hidden_dim):\n",
        "    super().__init__()    \n",
        "    self.attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "    self.norm1 = ResidualNorm(embed_dim)\n",
        "    self.feedforward = Feedforward(embed_dim, hidden_dim)\n",
        "    self.norm2 = ResidualNorm(embed_dim)\n",
        "\n",
        "  def forward(self, src_tokens, src_mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      src_tokens (n_batch, n_tokens, embed_dim): the source sequence\n",
        "      src_mask (n_batch, n_tokens): binary mask over the source\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): the encoder state\n",
        "    \"\"\"\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"EncoderBlock\")\n",
        "    ####################################################################\n",
        "    # First compute self-attention on the source tokens by passing them in\n",
        "    # as the queries, keys, and values to the attention module.\n",
        "    self_attention = ...\n",
        "    # Next compute the norm of the self-attention result with a residual\n",
        "    # connection from the source tokens\n",
        "    normed_attention = ...\n",
        "    # Pass the normed attention result through the feedforward component\n",
        "    ff_out = ...\n",
        "    # Finally compute the norm of the feedforward output with a residual\n",
        "    # connection from the normed attention output\n",
        "    out = ...\n",
        "    return out\n",
        "\n",
        "# Uncomment below to test your module\n",
        "# torch.manual_seed(522)\n",
        "# n_heads, batch_size, n_tokens, embed_dim, hidden_dim = 2, 1, 3, 4, 8\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# mask = torch.ones((batch_size, n_tokens))\n",
        "# encoder = EncoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "# encoder(tokens, mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o6BALWelvjG4"
      },
      "source": [
        "# to_remove solution\n",
        "class EncoderBlock(nn.Module):\n",
        "  def __init__(self, n_heads, embed_dim, hidden_dim):\n",
        "    super().__init__()    \n",
        "    self.attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "    self.norm1 = ResidualNorm(embed_dim)\n",
        "    self.feedforward = Feedforward(embed_dim, hidden_dim)\n",
        "    self.norm2 = ResidualNorm(embed_dim)\n",
        "\n",
        "  def forward(self, src_tokens, src_mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      src_tokens (n_batch, n_tokens, embed_dim): the source sequence\n",
        "      src_mask (n_batch, n_tokens): binary mask over the source\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): the encoder state\n",
        "    \"\"\"\n",
        "    self_attention = self.attention(src_tokens, src_tokens, src_tokens, src_mask)\n",
        "    normed_attention = self.norm1(self_attention, src_tokens)\n",
        "    ff_out = self.feedforward(normed_attention)\n",
        "    out = self.norm2(ff_out, normed_attention)\n",
        "    return out\n",
        "\n",
        "# torch.manual_seed(522)\n",
        "# n_heads, batch_size, n_tokens, embed_dim, hidden_dim = 2, 1, 3, 4, 8\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# mask = torch.ones((batch_size, n_tokens))\n",
        "# encoder = EncoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "# encoder(tokens, mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z-q1otdryMfO"
      },
      "source": [
        "If done correctly, you should see something like the following (numbers may vary but shape should be the same):\n",
        "```\n",
        "tensor([[[ 0.0051,  0.0022,  1.4105, -1.4179],\n",
        "         [-0.7053,  0.9854,  0.9762, -1.2564],\n",
        "         [-0.4003,  0.7551,  1.0888, -1.4436]]],\n",
        "       grad_fn=<NativeLayerNormBackward>)\n",
        "```\n",
        "\n",
        "[*Click for solution*](https://github.com/CIS-522/course-content/blob/main/tutorials/W10_NLP/solutions/W10_Tutorial2_Solution_Ex04.py)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7RaAaNO3DgMa"
      },
      "source": [
        "Now that we have our Encoder Block, we can chain these together in a stack to get the full Encoder module. We will include the embedding layer and positional encoding step of the source tokens here as well. The input to this module then will be a tensor of a batch of token IDs and corresponding mask.\n",
        "\n",
        "For instance, if our entire corpus was the English sentence: `Cat sat on the mat` and we tokenized by word, our vocab size would be 5 as there are 4 unique words. Converting this sentence to IDs would be `[[0,1,2,3,4]]`.\n",
        "\n",
        "The code for the Encoder module is provided below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HRgib6BjDuIo"
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "  def __init__(self, vocab_size, embed_dim, hidden_dim, n_heads, n_blocks):\n",
        "    super().__init__()\n",
        "    self.embedding = nn.Embedding(vocab_size, embed_dim)\n",
        "    self.positional_encoding = PositionalEncoder(embed_dim)\n",
        "    self.encoder_blocks = nn.ModuleList([\n",
        "      EncoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "      for _ in range(n_blocks)])\n",
        "  \n",
        "  def forward(self, src_tokens, src_mask):\n",
        "    x = self.embedding(src_tokens)\n",
        "    x = self.positional_encoding(x)\n",
        "    for block in self.encoder_blocks:\n",
        "      x = block(x, src_mask)\n",
        "    return x\n",
        "\n",
        "\n",
        "vocab_size = 5\n",
        "n_blocks, n_heads, batch_size, embed_dim, hidden_dim = 10, 2, 1, 4, 8\n",
        "enc = Encoder(vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "src_tokens = torch.IntTensor([[0,1,2,3,4]])\n",
        "src_mask = torch.IntTensor([[1,1,1,1,1]])\n",
        "enc(src_tokens, src_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZgRqX5_c8nbm"
      },
      "source": [
        "### The Decoder\n",
        "\n",
        "Like the encoder, the decoder is made up of a stack of repeating Decoder Blocks. Decoder Blocks are similar to the Encoder ones with an additional multi-head attention component that doesn't use self-attention, but instead gets the queries from the decoder's self-attention component and the keys and values from the encoder's output."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7GqPoroY9WEv"
      },
      "source": [
        "#### Exercise 5: Decoder Block Module\n",
        "In this exercise you will create the forward method of the PyTorch module representing the Decoder Block of the Transformer. The Decoder Block has the following architecture:\n",
        "1. a multi-head attention using self-attention\n",
        "2. 1st residual layer norm\n",
        "3. a 2nd multi-head attention that incorporates the encoder output\n",
        "4. 2nd residual layer norm\n",
        "5. feed-forward model\n",
        "6. 3rd residual layer norm"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J2C7z-h4Awqt"
      },
      "source": [
        "class DecoderBlock(nn.Module):\n",
        "  def __init__(self, n_heads, embed_dim, hidden_dim):\n",
        "    super().__init__()    \n",
        "    self.self_attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "    self.norm1 = ResidualNorm(embed_dim)\n",
        "    self.encoder_attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "    self.norm2 = ResidualNorm(embed_dim)\n",
        "    self.feedforward = Feedforward(embed_dim, hidden_dim)\n",
        "    self.norm3 = ResidualNorm(embed_dim)\n",
        "\n",
        "  def forward(self, tgt_tokens, tgt_mask, encoder_state, src_mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      tgt_tokens (n_batch, n_tokens, embed_dim): the target sequence\n",
        "      tgt_mask (n_batch, n_tokens): binary mask over the target tokens\n",
        "      encoder_state (n_batch, n_tokens, embed_dim): the output of the encoder pass\n",
        "      src_mask (n_batch, n_tokens): binary mask over the source tokens\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): the decoder state\n",
        "    \"\"\"\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"DecoderBlock\")\n",
        "    ####################################################################\n",
        "    # First compute self-attention on the target tokens by passing them in\n",
        "    # as the queries, keys, and values to the attention module along with the\n",
        "    # target mask.\n",
        "    self_attention = ...\n",
        "    # Next compute the norm of the self-attention result with a residual\n",
        "    # connection from the target tokens\n",
        "    normed_self_attention = ...\n",
        "    # Compute the encoder attention by using the normed self-attention output as\n",
        "    # the queries and the encoder state as the keys and values along with the\n",
        "    # source mask.\n",
        "    encoder_attention = ...\n",
        "    # Next compute the norm of the encoder attention result with a residual\n",
        "    # connection from the normed self-attention\n",
        "    normed_encoder_attention = ...\n",
        "    # Pass the normed encoder attention result through the feedforward component\n",
        "    ff_out = ...\n",
        "    # Finally compute the norm of the feedforward output with a residual\n",
        "    # connection from the normed attention output\n",
        "    out = ...\n",
        "    return out\n",
        "\n",
        "# Uncomment below to test your module\n",
        "# torch.manual_seed(522)\n",
        "# n_heads, batch_size, n_tokens, embed_dim, hidden_dim = 2, 1, 3, 4, 8\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# src_mask = torch.ones((batch_size, n_tokens))\n",
        "# tgt_mask = torch.ones((batch_size, n_tokens))\n",
        "# encoder = EncoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "# decoder = DecoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "# encoder_state = encoder(tokens, src_mask)\n",
        "# decoder(tokens, tgt_mask, encoder_state, src_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cpETDQkXvpVo"
      },
      "source": [
        "# to_remove solution\n",
        "class DecoderBlock(nn.Module):\n",
        "  def __init__(self, n_heads, embed_dim, hidden_dim):\n",
        "    super().__init__()    \n",
        "    self.self_attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "    self.norm1 = ResidualNorm(embed_dim)\n",
        "    self.encoder_attention = MultiHeadAttention(n_heads, embed_dim)\n",
        "    self.norm2 = ResidualNorm(embed_dim)\n",
        "    self.feedforward = Feedforward(embed_dim, hidden_dim)\n",
        "    self.norm3 = ResidualNorm(embed_dim)\n",
        "\n",
        "  def forward(self, tgt_tokens, tgt_mask, encoder_state, src_mask):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      tgt_tokens (n_batch, n_tokens, embed_dim): the target sequence\n",
        "      tgt_mask (n_batch, n_tokens): binary mask over the target tokens\n",
        "      encoder_state (n_batch, n_tokens, embed_dim): the output of the encoder pass\n",
        "      src_mask (n_batch, n_tokens): binary mask over the source tokens\n",
        "    Returns:\n",
        "      (n_batch, n_tokens, embed_dim): the decoder state\n",
        "    \"\"\"\n",
        "    self_attention = self.self_attention(tgt_tokens, tgt_tokens, tgt_tokens, tgt_mask)\n",
        "    normed_self_attention = self.norm1(self_attention, tgt_tokens)\n",
        "    encoder_attention = self.encoder_attention(normed_self_attention, encoder_state, encoder_state, src_mask)\n",
        "    normed_encoder_attention = self.norm2(encoder_attention, normed_self_attention)\n",
        "    ff_out = self.feedforward(normed_encoder_attention)\n",
        "    out = self.norm2(ff_out, normed_encoder_attention)\n",
        "    return out\n",
        "\n",
        "# torch.manual_seed(522)\n",
        "# n_heads, batch_size, n_tokens, embed_dim, hidden_dim = 2, 1, 3, 4, 8\n",
        "# tokens = torch.normal(0, 1, (batch_size, n_tokens, embed_dim))\n",
        "# src_mask = torch.ones((batch_size, n_tokens))\n",
        "# tgt_mask = torch.ones((batch_size, n_tokens))\n",
        "# encoder = EncoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "# decoder = DecoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "# encoder_state = encoder(tokens, src_mask)\n",
        "# decoder(tokens, tgt_mask, encoder_state, src_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCyIl3NnyQ-V"
      },
      "source": [
        "If done correctly, you should see something like the following (numbers may vary but shape should be the same):\n",
        "```\n",
        "tensor([[[ 1.0841,  0.3174,  0.2326, -1.6340],\n",
        "         [ 0.4667,  1.1922, -0.1277, -1.5312],\n",
        "         [ 0.6861,  0.9347,  0.0088, -1.6296]]],\n",
        "       grad_fn=<NativeLayerNormBackward>)\n",
        "```\n",
        "[*Click for solution*](https://github.com/CIS-522/course-content/blob/main/tutorials/W10_NLP/solutions/W10_Tutorial2_Solution_Ex05.py)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TmLB6F9QU17z"
      },
      "source": [
        "The Decoder module ends up just the same as the Encoder module with one key difference: the forward method needs to also accept the output of the encoder as well as the source token mask.\n",
        "\n",
        "For instance, let's assume we are doing a translation task and want to translate the English `Cat sat on the mat` into the French `Chat assis sur le tapis`. Out target vocab size is also 5 and would be similarly converted into IDs as `[[0,1,2,3,4]]`.\n",
        "\n",
        "The code for the Decoder module is presented below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ddhz5rGeUtvT"
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "  def __init__(self, vocab_size, embed_dim, hidden_dim, n_heads, n_blocks):\n",
        "    super().__init__()\n",
        "    self.embedding = nn.Embedding(vocab_size, embed_dim)\n",
        "    self.positional_encoding = PositionalEncoder(embed_dim)\n",
        "    self.decoder_blocks = nn.ModuleList([\n",
        "      DecoderBlock(n_heads, embed_dim, hidden_dim)\n",
        "      for _ in range(n_blocks)])\n",
        "  \n",
        "  def forward(self, tgt_tokens, tgt_mask, encoder_state, src_mask):\n",
        "    x = self.embedding(tgt_tokens)\n",
        "    x = self.positional_encoding(x)\n",
        "    for block in self.decoder_blocks:\n",
        "      x = block(x, tgt_mask, encoder_state, src_mask)\n",
        "    return x\n",
        "\n",
        "vocab_size = 5\n",
        "n_blocks, n_heads, batch_size, embed_dim, hidden_dim = 10, 2, 1, 4, 8\n",
        "tgt_tokens = torch.IntTensor([[0,1,2,3,4]])\n",
        "tgt_mask = src_mask = torch.IntTensor([[1,1,1,1,1]])\n",
        "enc_state = torch.randn((1,5,4))\n",
        "dec = Decoder(vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "dec(tgt_tokens, tgt_mask, enc_state, src_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vl3MHxUOCrt0"
      },
      "source": [
        "### The Full Transformer Model\n",
        "\n",
        "We can now put the Encoder and Decoder together to produce the full Transformer model.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vZjirF47WdHS"
      },
      "source": [
        "#### Exercise 6: Transformer Module\n",
        "\n",
        "In the last exercise for this section you will implement the forward method of the full Transformer module. First you will apply the source tokens and mask to the Encoder to get its output, then use that along with the target tokens and mask to produce the Decoder output. Finally we run the Decoder output through a linear layer to transform the embeddings back into vocab ID scores in order to determine the actual next word prediction."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D7A75izaW34X"
      },
      "source": [
        "class Transformer(nn.Module):\n",
        "  def __init__(self, src_vocab_size, tgt_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks):\n",
        "    super().__init__()\n",
        "    self.encoder = Encoder(src_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "    self.decoder = Decoder(tgt_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "    self.out = nn.Linear(embed_dim, tgt_vocab_size)\n",
        "\n",
        "  def forward(self, src_tokens, src_mask, tgt_tokens, tgt_mask):\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"DecoderBlock\")\n",
        "    ####################################################################\n",
        "    # Compute the encoder output state from the source tokens and mask\n",
        "    encoder_state = ...\n",
        "    # Compute the decoder output state from the target tokens and mask as well\n",
        "    # as the encoder state and source mask\n",
        "    decoder_state = ...\n",
        "    # Compute the vocab scores by passing the decoder state through the output\n",
        "    # linear layer\n",
        "    out = ...\n",
        "    return out\n",
        "\n",
        "# Uncomment below to test your module\n",
        "# torch.manual_seed(522)\n",
        "# src_vocab_size = tgt_vocab_size = 5\n",
        "# n_blocks, n_heads, batch_size, embed_dim, hidden_dim = 10, 2, 1, 4, 8\n",
        "# src_tokens = tgt_tokens = torch.IntTensor([[0,1,2,3,4]])\n",
        "# src_mask = tgt_mask = torch.IntTensor([[1,1,1,1,1]])\n",
        "\n",
        "# transformer = Transformer(src_vocab_size, tgt_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "# transformer(src_tokens, src_mask, tgt_tokens, tgt_mask)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ajo0WSHqjtAq"
      },
      "source": [
        "# to_remove solution\n",
        "class Transformer(nn.Module):\n",
        "  def __init__(self, src_vocab_size, tgt_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks):\n",
        "    super().__init__()\n",
        "    self.encoder = Encoder(src_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "    self.decoder = Decoder(tgt_vocab_size, embed_dim, hidden_dim, n_heads, n_blocks)\n",
        "    self.out = nn.Linear(embed_dim, tgt_vocab_size)\n",
        "\n",
        "  def forward(self, src_tokens, src_mask, tgt_tokens, tgt_mask):\n",
        "    # Compute the encoder output state from the source tokens and mask\n",
        "    encoder_state = self.encoder(src_tokens, src_mask)\n",
        "    # Compute the decoder output state from the target tokens and mask as well\n",
        "    # as the encoder state and source mask\n",
        "    decoder_state = self.decoder(tgt_tokens, tgt_mask, encoder_state, src_mask)\n",
        "    # Compute the vocab scores by passing the decoder state through the output\n",
        "    # linear layer\n",
        "    out = self.out(decoder_state)\n",
        "    return out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24ZGB98DbdPW"
      },
      "source": [
        "If done correctly, you should see something like the following (numbers may vary but shape should be the same):\n",
        "```\n",
        "tensor([[[-0.1359,  0.5821, -0.5340, -0.7582,  0.0687],\n",
        "         [ 0.0085,  0.1495, -0.1809, -0.9419, -0.4447],\n",
        "         [-0.5151,  0.5056,  0.8117,  0.0047, -0.6726],\n",
        "         [-0.1393, -0.2927,  0.9524, -0.5759, -1.3004],\n",
        "         [-0.4090,  1.4626,  0.2387, -0.1716, -0.2155]]],\n",
        "       grad_fn=<AddBackward0>)\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kz0tL-xaBLkv"
      },
      "source": [
        "### Summary\n",
        "We've covered all the building blocks that make up the Transformer network architecture from the attention mechanism up to the fully combined encoder-decoder setup. The module versions presented here were often simplified in some ways and made more verbose in others to emphasize what each component is doing."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "SblV7fWsC4jx"
      },
      "source": [
        "#@title Video : Transformer Architecture\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"_sKZpAptIZk\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "guMptHSvoaIz"
      },
      "source": [
        "*Estimated time: 85 minutes since start*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G31Vd2hVzuQt"
      },
      "source": [
        "---\n",
        "##Section 2: BERT and friends"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_4b-wW01AUQ8",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Bert and Friends\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"gEkmPb0140w\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IwJPWkxHj8a8"
      },
      "source": [
        "---\n",
        "## Section 3: BERT\n",
        "\n",
        "BERT, or Bidrectional Encoder Representations from Transforms, is a Transformer-based machine learning technique for NLP pre-training developed by Google. The original English BERT has two models: \n",
        "\n",
        "1. BERT$_{BASE}$: $12$ encoders with $12$ bidirectional self-attention heads\n",
        "2. BERT$_{LARGE}$: $24$ encoders with $24$ bidirectional self-attention heads \n",
        "\n",
        "Both models are pre-trained with unlabeled data extracted from BooksCorpus with $800$M words and Wikipedia with $2500$M words. Importantly, unlike context-free models like GloVe or word2vec, BERT takes context into account for each occurrence of a given word.  For instance, whereas the vector for \"running\" will have the same word2vec vector representation for both of its occurrences in the sentences \"He is running a company\" and \"He is running a marathon\", BERT will provide a contextualized embedding that will be different according to the sentence.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lkGd01-Wsthf",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Using BERT\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"sFQGoswoaeI\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kRmL_Thxoi46"
      },
      "source": [
        "*Estimated time: 95 minutes since start*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jlRZjEpPtMAm"
      },
      "source": [
        "---\n",
        "## Section 4: RoBERTa"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CJ2X4VP0ssNE"
      },
      "source": [
        "As can be seen by the name, RoBERTa builds on BERT, modifying key hyperparameters. It removes the next-sentence pretraining objective and trains with much larger mini-batches and learning rates. \n",
        "\n",
        "Spend some time playing with RoBERTa natural language inference at https://huggingface.co/roberta-large-mnli?text=The+dog+was+lost.+I+lost+an+animal.+. Additionally, spend some time looking at the other model examples at https://github.com/huggingface/transformers.\n",
        "\n",
        "\n",
        "Roberta, introduced [here](https://arxiv.org/abs/1907.11692), stands for A Robustly Optimized BERT Pretraining Approach. RoBERTa builds on BERT's language masking strategy, where the system learns to predict intentionally hidden sections of text within otherwise unannotated language samples. It modifies key hyperparameters in BERT, removing BERT's next-sentence pretraining objective. This, alongside the larger mini-batch and learning rate, allows RoBERTA to improve ont he masked language modeling objective compared with BERT and leads to better downstream task performance. The changes in training causing a substantial increase in performance leads to the believe that BERt was relatively undertrained. \n",
        "\n",
        "\n",
        "## TweetEVAL Sentiment Analysis\n",
        "\n",
        "We utilize an already finetuned network on sentiment analysis to write an example for sentiment analysis. This is a roBERTa model trained on ~58 million tweets and finetuned for sentiment analysis with the [TweetEval](https://arxiv.org/pdf/2010.12421.pdf) benchmark. We use huggingface to implement this. \n",
        "\n",
        "\n",
        "First, we preprocess the text and download our fine-tuned model. By fine-tuning, this model took the initial training done in the roBERTa paper and trained it additionally on a downstream sentiment analysis task. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X7VdwufAsC8g"
      },
      "source": [
        "def preprocess(text):\n",
        "    new_text = []\n",
        "\n",
        "\n",
        "    for t in text.split(\" \"):\n",
        "        t = '@user' if t.startswith('@') and len(t) > 1 else t\n",
        "        t = 'http' if t.startswith('http') else t\n",
        "        new_text.append(t)\n",
        "    return \" \".join(new_text)\n",
        "\n",
        "# Tasks:\n",
        "# emoji, emotion, hate, irony, offensive, sentiment\n",
        "# stance/abortion, stance/atheism, stance/climate, stance/feminist, stance/hillary\n",
        "\n",
        "task='sentiment'\n",
        "MODEL = f\"cardiffnlp/twitter-roberta-base-{task}\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
        "\n",
        "# download label mapping\n",
        "labels=[]\n",
        "mapping_link = f\"https://raw.githubusercontent.com/cardiffnlp/tweeteval/main/datasets/{task}/mapping.txt\"\n",
        "with urllib.request.urlopen(mapping_link) as f:\n",
        "    html = f.read().decode('utf-8').split(\"\\n\")\n",
        "    csvreader = csv.reader(html, delimiter='\\t')\n",
        "labels = [row[1] for row in csvreader if len(row) > 1]\n",
        "\n",
        "\n",
        "model = AutoModelForSequenceClassification.from_pretrained(MODEL)\n",
        "model.save_pretrained(MODEL)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9BO-LGbisIFt"
      },
      "source": [
        "Now that we have our pytorch model ready, let's play with some text. As seen below, we take some input text (feel free to change this to whatever text you want to see the sentiment). The text is preprocessed usingour function above, then tokenized using the huggingface autotokenizer. Then, we sort the sentiments by their relative probabilities, which is what we see at the end. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UNHjqYEMsJlp"
      },
      "source": [
        "text = \"I'm sad :(\"\n",
        "text = preprocess(text)\n",
        "encoded_input = tokenizer(text, return_tensors='pt')\n",
        "output = model(**encoded_input)\n",
        "scores = output[0][0].detach().numpy()\n",
        "scores = softmax(scores)\n",
        "\n",
        "\n",
        "\n",
        "ranking = np.argsort(scores)\n",
        "ranking = ranking[::-1]\n",
        "for i in range(scores.shape[0]):\n",
        "    l = labels[ranking[i]]\n",
        "    s = scores[ranking[i]]\n",
        "    print(f\"{i+1}) {l} {np.round(float(s), 4)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bP1rPKAOooDd"
      },
      "source": [
        "*Estimated time: 110 minutes since start*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WEgqsbk5zy0J"
      },
      "source": [
        "---\n",
        "## Section 5: BERT variations (DistilBERT)\n",
        "\n",
        "https://arxiv.org/abs/1910.01108\n",
        "\n",
        "\n",
        "DistilBERT, as the name suggests, is a \"distilled\" version of BERT: smaller, faster, cheaper, and lighter. Often times, having very large models is infeasible, as it requires a lot of compute time and resources. Specifically, we often need to run models on smaller devices, without the ability to run many large GPUs to train. DistilBERT is a pre-trained general-purpose language representation model, which we can then fine-tune to achieve good performance on a number of tasks. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6vJP1ZEaDXbO"
      },
      "source": [
        "Let's use DistilBERT to write a small question answering system. Question answering systems automatically respond to a given query. The input will be framed with context and the question. For example:\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "Context :\n",
        "The US has passed the peak on new coronavirus cases, President Donald Trump said and predicted that some states would reopen this month. The US has over 637,000 confirmed Covid-19 cases and over 30,826 deaths, the highest for any country in the world. \\\\\n",
        "Question:\n",
        "What was President Donald Trump's prediction?\n",
        "\n",
        "\n",
        "---\n",
        "Answer: \n",
        "some states would reopen this month.\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vA4C_ZpPD7PD"
      },
      "source": [
        "from transformers import DistilBertTokenizer, DistilBertForQuestionAnswering\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased',return_token_type_ids = True)\n",
        "model = DistilBertForQuestionAnswering.from_pretrained('distilbert-base-uncased-distilled-squad', return_dict=False)\n",
        "\n",
        "context = \"The US has passed the peak on new coronavirus cases, \" \\\n",
        "          \"President Donald Trump said and predicted that some states would reopen this month. \" \\\n",
        "          \"The US has over 637,000 confirmed Covid-19 cases and over 30,826 deaths, the highest for any country in the world.\"\n",
        "\n",
        "question = \"What was President Donald Trump's prediction?\"\n",
        "\n",
        "encoding = tokenizer.encode_plus(question, context)\n",
        "\n",
        "\n",
        "input_ids, attention_mask = encoding[\"input_ids\"], encoding[\"attention_mask\"]\n",
        "\n",
        "start_scores, end_scores = model(torch.tensor([input_ids]), attention_mask=torch.tensor([attention_mask]))\n",
        "\n",
        "ans_tokens = input_ids[torch.argmax(start_scores) : torch.argmax(end_scores)+1]\n",
        "answer_tokens = tokenizer.convert_ids_to_tokens(ans_tokens , skip_special_tokens=True)\n",
        "\n",
        "print (\"\\nQuestion: \",question)\n",
        "print (\"\\nAnswer Tokens: \")\n",
        "print (answer_tokens)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "answer_tokens_to_string = tokenizer.convert_tokens_to_string(answer_tokens)\n",
        "\n",
        "print (\"\\nAnswer : \",answer_tokens_to_string)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkLAUmkbEyeK"
      },
      "source": [
        "Cool! Go ahead and try your own questions and see how DistilBERT answers it! Let's try multiple questions at once (in a batch). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rxjk7g7iE8lT"
      },
      "source": [
        "\n",
        "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased',return_token_type_ids = True)\n",
        "model = DistilBertForQuestionAnswering.from_pretrained('distilbert-base-uncased-distilled-squad', return_dict = False)\n",
        "\n",
        "context = \"The US has passed the peak on new coronavirus cases, \" \\\n",
        "          \"President Donald Trump said and predicted that some states would reopen this month.\" \\\n",
        "          \"The US has over 637,000 confirmed Covid-19 cases and over 30,826 deaths, \" \\\n",
        "          \"the highest for any country in the world.\"\n",
        "\n",
        "print (\"\\n\\nContext : \",context)\n",
        "\n",
        "questions = [\"What was President Donald Trump's prediction?\",\n",
        "             \"How many deaths have been reported from the virus?\",\n",
        "             \"How many cases have been reported in the United States?\"]\n",
        "\n",
        "question_context_for_batch = []\n",
        "\n",
        "for question in questions :\n",
        "    question_context_for_batch.append((question, context))\n",
        "\n",
        "encoding = tokenizer.batch_encode_plus(question_context_for_batch,pad_to_max_length=True, return_tensors=\"pt\")\n",
        "input_ids, attention_mask = encoding[\"input_ids\"], encoding[\"attention_mask\"]\n",
        "start_scores, end_scores = model(input_ids, attention_mask=attention_mask)\n",
        "\n",
        "for index,(start_score,end_score,input_id) in enumerate(zip(start_scores,end_scores,input_ids)):\n",
        "    max_startscore = torch.argmax(start_score)\n",
        "    max_endscore = torch.argmax(end_score)\n",
        "    ans_tokens = input_ids[index][max_startscore: max_endscore + 1]\n",
        "    answer_tokens = tokenizer.convert_ids_to_tokens(ans_tokens, skip_special_tokens=True)\n",
        "    answer_tokens_to_string = tokenizer.convert_tokens_to_string(answer_tokens)\n",
        "    print (\"\\nQuestion: \",questions[index])\n",
        "    print (\"Answer: \", answer_tokens_to_string)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qoyEwsPPozXz"
      },
      "source": [
        "*Estimated time: 120 minutes since start*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yUmSToDB3bFH"
      },
      "source": [
        "---\n",
        "## Section 6: Explaining language models\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A2EiG7Rs3kvb",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Explaining language models\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"G38ZZNnXaQs\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_aKo3ReY_9sT"
      },
      "source": [
        "#### Questions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5mkr7X0a__ko",
        "cellView": "form"
      },
      "source": [
        "#@markdown Why would you expect part of speech tagging to be done closer to the input, and co-reference to be done more deeply in the network?\n",
        "#report to Airtable\n",
        "NLP_network_structure = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PzJOIertAVq9",
        "cellView": "form"
      },
      "source": [
        "#@markdown Why are byte pair encodings problematic for using \"feature importance\" to understand what words \"cause\" a model to make a given prediction?\n",
        "BPE_interpretation = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C9cTkXEsAuV4",
        "cellView": "form"
      },
      "source": [
        "#@markdown Attention turns out not to be a great way of finding the most important words used by a model. Why not? (Hint: where might attention focus on the sentence: \"The movie was long and boring.\"?)\n",
        "interpreting_attention = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SdrEjN3n31Wx"
      },
      "source": [
        "There are lots of tools out there to help visualize what's going on in NLP systems. If you want (this is not an assignment), play around with the demos at https://pair-code.github.io/lit/demos/.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I9lghzopo3LF"
      },
      "source": [
        "*Estimated time: 130 minutes since start*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "52dVCGNMz9w-"
      },
      "source": [
        "---\n",
        "## Section 7: Bias in Embeddings\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6AU4_g3Ddoli",
        "cellView": "form"
      },
      "source": [
        "#@title Video : Bias in Embeddings\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"meUnCri_52c\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NfmYlSc_0BQz"
      },
      "source": [
        "You just saw how training on large amounts of historical text can introduce undesirable associations and outcomes. In this section we are going to explore this idea further as it pertains to coreference resolution.\n",
        "\n",
        "Coreference resolution is the NLP task of finding all the terms that refer to an entity in a passage of text e.g. what noun phrase does a pronoun refer to. This can be quite difficult, even for humans, if the passage is ambiguous enough.\n",
        "\n",
        "For example, in the sentence:\n",
        "\n",
        "`The customer asked to speak with the manager because he wanted to fix the billing error quickly.`\n",
        "\n",
        "what does `he` refer to? We can reasonably assume given the context that `he` refers to the customer. Furthermore, it shouldn't matter which pronoun (he/she/they) was in that spot, it should still refer back to the customer.\n",
        "\n",
        "However this is not the case with some models! For example, here is the output of Huggingface's [Neural Coreference model](https://github.com/huggingface/neuralcoref) when we use `she` as the pronoun:\n",
        "\n",
        "![coref with she](https://imgur.com/4FmN3ZC.png)\n",
        "\n",
        "You can see that `she` is scored against all detected noun phrases and gets the highest score with `the customer`. So far so good. Now let's try it with `he` instead:\n",
        "\n",
        "![coref with he](https://imgur.com/SOyLrtg.png)\n",
        "\n",
        "The model has instead associated `he` with `the manager`, and quite strongly at that, even though that doesn't make sense contextually. As this is a neural-based model trained on historical data, one possibility is there were many instances where `he` and `manager` were associated, enough to get \"distracted\" by that signal over the rest of the sentence.\n",
        "\n",
        "\n",
        "\n",
        "As was mentioned in the video, many people are actively working toward both identifying and mitigating these undesirable behaviors and also educating researchers, practitioners, and the general public about these issues. For instance, the sample sentence used above was taken from the [Winogender Schemas](https://github.com/rudinger/winogender-schemas), a set of sample sentences to check the variance in outcomes when only a single pronoun is changed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ftnc8GGUtL7O"
      },
      "source": [
        "#### Exercise 6: Explore Bias in Coreference Resolution Models\n",
        "\n",
        "Two different coreference resolution models that have nice online demos are from [Huggingface](https://huggingface.co/coref/) and [AllenNLP](https://demo.allennlp.org/coreference-resolution). In this exercise, you will explore a variety of sentences with these two tools and see how they compare. Try the following sentences in both and see how they handle the change in pronoun:\n",
        "\n",
        "`The doctor berated the nurse. He had come in late for the meeting.`\n",
        "\n",
        "`The doctor berated the nurse. She had come in late for the meeting.`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ClZxGxhJBZa0"
      },
      "source": [
        "#### Questions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aUAhscXBBdMX",
        "cellView": "form"
      },
      "source": [
        "#@markdown Did Huggingface get it right?\n",
        "huggingface_bias = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DAsS-ChpBuQv",
        "cellView": "form"
      },
      "source": [
        "#@markdown Did Allen Institute get it right?\n",
        "#report to Airtable\n",
        "allenInst_bias = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VKW9WhInB5GM",
        "cellView": "form"
      },
      "source": [
        "#@markdown How might you \"fine tune\" Bert to reduce such errors?\n",
        "#report to Airtable\n",
        "fine_tune_away_bias = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gbgTqgB5rH_W"
      },
      "source": [
        "---\n",
        "# Wrap up\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "98YXpSntrY6k",
        "cellView": "form"
      },
      "source": [
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "\n",
        "import time\n",
        "import numpy as np\n",
        "import urllib.parse\n",
        "from IPython.display import IFrame\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefill_fields = {}\n",
        "  for key in fields:\n",
        "      new_key = 'prefill_' + key\n",
        "      prefill_fields[new_key] = fields[key]\n",
        "  prefills = urllib.parse.urlencode(prefill_fields)\n",
        "  src = src + prefills\n",
        "  return src\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "try: t4;\n",
        "except NameError: t4 = time.time()\n",
        "try: t5;\n",
        "except NameError: t5 = time.time()\n",
        "try: t6;\n",
        "except NameError: t6 = time.time()\n",
        "try: t7;\n",
        "except NameError: t7 = time.time()\n",
        "\n",
        "# autofill fields if they are not present\n",
        "# a missing pennkey and pod will result in an Airtable warning\n",
        "# which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: NLP_network_structure;\n",
        "except NameError: NLP_network_structure = \"\" \n",
        "try: BPE_interpretation;\n",
        "except NameError: BPE_interpretation = \"\" \n",
        "try: interpreting_attention;\n",
        "except NameError: interpreting_attention = \"\" \n",
        "try: huggingface_bias;\n",
        "except NameError: huggingface_bias = \"\" \n",
        "try: allenInst_bias;\n",
        "except NameError: allenInst_bias = \"\" \n",
        "try: fine_tune_away_bias;\n",
        "except NameError: fine_tune_away_bias = \"\" \n",
        "\n",
        "times = np.array([t1,t2,t3,t4,t5,t6,t7])-t0\n",
        "\n",
        "fields = {\n",
        "    \"my_pennkey\": my_pennkey,\n",
        "    \"my_pod\": my_pod, \n",
        "    \"NLP_network_structure\": NLP_network_structure,\n",
        "    \"BPE_interpretation\": BPE_interpretation,\n",
        "    \"interpreting_attention\": interpreting_attention,\n",
        "    \"huggingface_bias\": huggingface_bias,\n",
        "    \"allenInst_bias\": allenInst_bias,\n",
        "    \"fine_tune_away_bias\": fine_tune_away_bias,\n",
        "    \"cumulative_times\": times\n",
        "}\n",
        "\n",
        "src = \"https://airtable.com/embed/shrfeQ4zBWMSZSheB?\"\n",
        "\n",
        "display(IFrame(src = prefill_form(src, fields), width = 800, height = 400))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oDE0MJbb5dLH"
      },
      "source": [
        "## Feedback\n",
        "How could this session have been better? How happy are you in your group? How do you feel right now?\n",
        "\n",
        "Feel free to use the embeded form below or use this link:\n",
        "<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://airtable.com/shrNSJ5ECXhNhsYss\">https://airtable.com/shrNSJ5ECXhNhsYss</a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IPPjyA-H5kLE"
      },
      "source": [
        "display(IFrame(src=\"https://airtable.com/embed/shrNSJ5ECXhNhsYss?backgroundColor=red\", width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}